import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from scipy.stats import zscore
from DatabaseReader import DatabaseReader


def plot_time_distribution():
    """绘制历史订单数据中完成时间的分布图"""
    # 获取数据
    db_reader = DatabaseReader()
    df = db_reader.read_robot_order()
    print(df.shape)
    # 提取 taketime 数据并去除缺失值
    taketime = df['taketime'].dropna()

    # ---- 去掉 taketime 小于 0 的数据 ----
    taketime = taketime[taketime >= 0]

    # ---- 异常值去除 ----
    # 使用 Z-score 去除异常值，通常认为 Z-score > 3 或 < -3 的数据为异常值
    z_scores = zscore(taketime)
    taketime_clean = taketime[(z_scores > -3) & (z_scores < 3)]
    print(taketime_clean.shape)

    # 设置绘图样式
    # sns.set(style="whitegrid")

    # 创建图形对象
    fig, ax1 = plt.subplots(figsize=(10, 6))

    # ---- 绘制 PDF (概率密度函数) ----
    # 设置 RGB 颜色：例如：红色 (1, 0, 0) 和蓝色 (0, 0, 1)
    color_red = (223/255, 122/255, 94/255)
    color_green = (130/255, 178/255, 154/255)
    color_yellow = (242/255, 204/255, 142/255)
    sns.kdeplot(taketime_clean, ax=ax1, color=color_yellow, shade=True, label='PDF', linewidth=2, bw_adjust=0.5)  # 红色

    # ---- 绘制 CDF (累积分布函数) ----
    ax2 = ax1.twinx()  # 创建一个共享 x 轴的第二个 y 轴
    sns.ecdfplot(taketime_clean, ax=ax2, color=color_green, label='CDF', linewidth=2)  # 绿色

    # 设置标题和标签
    ax1.set_title('The distribution of travel time')
    ax1.set_xlabel('Taketime (seconds)')
    ax1.set_ylabel('Density', color=color_yellow)  # 红色
    ax2.set_ylabel('CDF', color=color_green)  # 绿色

    ax1.set_xlim(left=0)

    # 设置颜色标签
    ax1.legend(loc='upper left')
    ax2.legend(loc='upper right')

    # 调整图形布局
    plt.tight_layout()

    # 保存图像为PDF文件
    # plt.savefig('taketime_analysis_clean_colored.pdf', dpi=300)

    # 显示图像
    plt.show()


def plot_length_distribution():
    """绘制历史订单数据中路径长度的分布图"""
    # 获取数据
    db_reader = DatabaseReader()
    df = db_reader.read_robot_order()

    # 提取 taketime 数据并去除缺失值
    taketime = df['takelength'].dropna()

    # ---- 去掉 taketime 小于 0 的数据 ----
    taketime = taketime[taketime >= 0]

    # ---- 异常值去除 ----
    # 使用 Z-score 去除异常值，通常认为 Z-score > 3 或 < -3 的数据为异常值
    z_scores = zscore(taketime)
    taketime_clean = taketime[(z_scores > -3) & (z_scores < 3)]

    # 设置绘图样式
    # sns.set(style="whitegrid")

    # 创建图形对象
    fig, ax1 = plt.subplots(figsize=(10, 6))

    # ---- 绘制 PDF (概率密度函数) ----
    # 设置 RGB 颜色：例如：红色 (1, 0, 0) 和蓝色 (0, 0, 1)
    color_red = (223/255, 122/255, 94/255)
    color_green = (130/255, 178/255, 154/255)
    color_yellow = (242/255, 204/255, 142/255)
    sns.kdeplot(taketime_clean, ax=ax1, color=color_yellow, shade=True, label='PDF', linewidth=2, bw_adjust=0.5)  # 红色

    # ---- 绘制 CDF (累积分布函数) ----
    ax2 = ax1.twinx()  # 创建一个共享 x 轴的第二个 y 轴
    sns.ecdfplot(taketime_clean, ax=ax2, color=color_green, label='CDF', linewidth=2)  # 绿色

    # 设置标题和标签
    ax1.set_title('The distribution of travel time')
    ax1.set_xlabel('Taketime (seconds)')
    ax1.set_ylabel('Density', color=color_yellow)  # 红色
    ax2.set_ylabel('CDF', color=color_green)  # 绿色

    ax1.set_xlim(left=0)

    # 设置颜色标签
    ax1.legend(loc='upper left')
    ax2.legend(loc='upper right')

    # 调整图形布局
    plt.tight_layout()

    # 保存图像为PDF文件
    # plt.savefig('taketime_analysis_clean_colored.pdf', dpi=300)

    # 显示图像
    plt.show()


if __name__ == '__main__':
    # test
    plot_length_distribution()
